#include <iostream>
#include "cutlass/layout/layout.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/coord.h"
#include "cutlass/semaphore.h"

#include <cuda_runtime_api.h>
#include <cooperative_groups.h>

using namespace std;
using namespace cutlass;


__global__ void semaphore_test(int* lock)
{
    auto cg = cooperative_groups::this_thread_block();

    Semaphore a(lock, threadIdx.x);

    a.release(7);

    printf("thread_rank: %d, state: %d\n", cg.thread_rank(), a.get_state());

    a.wait(7);

    printf("after thread_rank: %d, state: %d\n", cg.thread_rank(), a.get_state());
}

int main(int argc, char* argv[])
{
    using layout = cutlass::layout::RowMajor;
    using Stride = layout::Stride;

    Stride stride = cutlass::make_Coord<layout::LongIndex>(512);
    layout l(stride);

    cutlass::MatrixCoord coord = cutlass::make_Coord<>(3, 2);
    layout::LongIndex index = l(coord);

    cout << "long index: " << index << endl;

    layout::LongIndex offset = 102897;
    MatrixCoord cord = l.inverse(offset);
    cout << "coord: " << cord.row() << ", " << cord.column() << endl;

    int* device_lock;

    cudaMalloc(&device_lock, sizeof(int));
    cudaMemset(device_lock, 0, sizeof(int));

    semaphore_test<<<1, 128>>>(device_lock);
    cudaDeviceSynchronize();
    return 0;
}